#!/usr/bin/env python3
"""
Redis-rs2 集群压力测试
测试集群在高并发下的表现
"""

import socket
import time
import threading
import sys
from concurrent.futures import ThreadPoolExecutor
import statistics

def create_resp_command(command, *args):
    """创建RESP格式的命令"""
    parts = [command] + list(args)
    resp = f"*{len(parts)}\r\n"
    for part in parts:
        resp += f"${len(str(part))}\r\n{part}\r\n"
    return resp.encode()

def send_command(sock, command_bytes):
    """发送命令并接收响应"""
    try:
        sock.send(command_bytes)
        response = sock.recv(1024)
        return response.decode('utf-8', errors='ignore')
    except Exception as e:
        return f"Error: {e}"

def worker_thread(worker_id, host, port, operations_per_thread, results):
    """工作线程函数"""
    try:
        # 创建连接
        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        sock.settimeout(10)
        sock.connect((host, port))
        
        success_count = 0
        error_count = 0
        latencies = []
        
        # 执行操作
        for i in range(operations_per_thread):
            start_time = time.time()
            
            # SET操作
            key = f"stress_key_{worker_id}_{i}"
            value = f"stress_value_{worker_id}_{i}"
            cmd = create_resp_command("SET", key, value)
            response = send_command(sock, cmd)
            
            if response and "+OK" in response:
                success_count += 1
                latency = (time.time() - start_time) * 1000  # 毫秒
                latencies.append(latency)
            else:
                error_count += 1
        
        sock.close()
        
        results[worker_id] = {
            'success': success_count,
            'errors': error_count,
            'latencies': latencies
        }
        
    except Exception as e:
        results[worker_id] = {
            'success': 0,
            'errors': operations_per_thread,
            'latencies': [],
            'connection_error': str(e)
        }

def cluster_stress_test(host='127.0.0.1', port=7001, 
                       threads=10, operations_per_thread=100):
    """集群压力测试"""
    print(f"🚀 Redis-rs2 集群压力测试")
    print(f"🎯 目标: {host}:{port}")
    print(f"🔥 并发线程: {threads}")
    print(f"📊 每线程操作数: {operations_per_thread}")
    print(f"📈 总操作数: {threads * operations_per_thread}")
    print("=" * 60)
    
    # 结果存储
    results = {}
    
    # 开始时间
    start_time = time.time()
    
    # 创建线程池
    with ThreadPoolExecutor(max_workers=threads) as executor:
        futures = []
        for i in range(threads):
            future = executor.submit(
                worker_thread, i, host, port, operations_per_thread, results
            )
            futures.append(future)
        
        # 等待所有线程完成
        for future in futures:
            future.result()
    
    # 结束时间
    end_time = time.time()
    total_time = end_time - start_time
    
    # 统计结果
    total_success = 0
    total_errors = 0
    all_latencies = []
    connection_errors = 0
    
    for worker_id, result in results.items():
        total_success += result['success']
        total_errors += result['errors']
        all_latencies.extend(result['latencies'])
        
        if 'connection_error' in result:
            connection_errors += 1
            print(f"❌ 线程 {worker_id} 连接失败: {result['connection_error']}")
    
    # 计算性能指标
    total_operations = total_success + total_errors
    success_rate = (total_success / total_operations) * 100 if total_operations > 0 else 0
    ops_per_second = total_success / total_time if total_time > 0 else 0
    
    # 延迟统计
    avg_latency = statistics.mean(all_latencies) if all_latencies else 0
    p95_latency = statistics.quantiles(all_latencies, n=20)[18] if len(all_latencies) >= 20 else 0
    p99_latency = statistics.quantiles(all_latencies, n=100)[98] if len(all_latencies) >= 100 else 0
    
    # 打印结果
    print("\n📊 测试结果:")
    print(f"   ⏱️  总耗时: {total_time:.2f} 秒")
    print(f"   ✅ 成功操作: {total_success}")
    print(f"   ❌ 失败操作: {total_errors}")
    print(f"   🔌 连接错误: {connection_errors}")
    print(f"   📈 成功率: {success_rate:.2f}%")
    print(f"   🚀 吞吐量: {ops_per_second:.0f} ops/sec")
    
    if all_latencies:
        print(f"\n⏱️  延迟统计:")
        print(f"   📊 平均延迟: {avg_latency:.2f} ms")
        print(f"   📊 P95延迟: {p95_latency:.2f} ms")
        print(f"   📊 P99延迟: {p99_latency:.2f} ms")
    
    print("\n" + "=" * 60)
    
    # 判断测试是否通过
    if success_rate >= 95 and ops_per_second >= 1000:
        print("✅ 集群压力测试通过")
        return True
    else:
        print("❌ 集群压力测试失败")
        return False

def main():
    host = sys.argv[1] if len(sys.argv) > 1 else '127.0.0.1'
    port = int(sys.argv[2]) if len(sys.argv) > 2 else 7001
    threads = int(sys.argv[3]) if len(sys.argv) > 3 else 10
    ops = int(sys.argv[4]) if len(sys.argv) > 4 else 100
    
    success = cluster_stress_test(host, port, threads, ops)
    sys.exit(0 if success else 1)

if __name__ == "__main__":
    main()